import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm
import warnings

warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


class DroughtEmergencyAnalyzer:
    """问题3：旱灾应急分析模型"""

    def __init__(self, problem2_result):
        # 基础参数（复用问题2配置）
        self.farm_size = 100  # 农场为100m×100m正方形
        self.min_soil_moisture = 0.22  # 作物存活最低湿度
        self.drought_river_ratio = 0.8  # 旱灾时河流供水量为问题2的80%
        self.emergency_coverage = 50  # 应急储水覆盖半径(米)

        # 作物参数（表1数据）
        self.crops = {
            '高粱': {'area': 0.5, 'survival_hum': 0.22, 'growth_hum': 0.27},  # 存活湿度+0.05
            '玉米': {'area': 0.3, 'survival_hum': 0.22, 'growth_hum': 0.28},  # 存活湿度+0.06
            '大豆': {'area': 0.2, 'survival_hum': 0.22, 'growth_hum': 0.26}  # 存活湿度+0.04
        }

        # 从问题2结果中获取相关参数
        self.pipe_daily_flow = problem2_result['管道日流量']  # 问题2设计的日流量(L)
        self.tank_total_volume = problem2_result['储水罐容积']  # 储水罐总容积(L)
        self.sprinkler_positions = problem2_result['喷头位置']  # 喷头坐标(m)
        self.daily_water_demand = problem2_result['日均需水量']  # 问题2中7月日均需水量(L)

        # 应急储备分析参数
        self.drought_probabilities = [10, 30, 50, 80, 100]  # 旱灾概率(%)

    def calculate_drought_supply(self):
        """计算旱灾期间的供水能力"""
        # 1. 河流供水量（80%设计流量）
        river_supply = self.pipe_daily_flow * self.drought_river_ratio

        # 2. 应急储水可利用量（假设储备比例r）
        # 先按30%储备计算基础值，后续会迭代优化
        emergency_ratio = 0.3
        emergency_supply = self.tank_total_volume * emergency_ratio

        # 3. 总供水量
        total_supply = river_supply + emergency_supply
        print(f"旱灾期间日供水量：河流{river_supply:.0f}L + 应急{emergency_supply:.0f}L = 总计{total_supply:.0f}L")
        print(f"日均需水量：{self.daily_water_demand:.0f}L")

        return {
            'river_supply': river_supply,
            'emergency_supply': emergency_supply,
            'total_supply': total_supply
        }

    def analyze_crop_survival(self, supply_data):
        """分析旱灾时作物存活与正常生长面积"""
        # 1. 供水缺口计算
        supply_gap = self.daily_water_demand - supply_data['total_supply']
        is_shortage = supply_gap > 0

        # 2. 按作物优先级分配水资源（高粱>玉米>大豆，按种植面积比例）
        crop_ratios = {
            '高粱': self.crops['高粱']['area'] / 1.0,
            '玉米': self.crops['玉米']['area'] / 1.0,
            '大豆': self.crops['大豆']['area'] / 1.0
        }

        # 3. 计算实际覆盖面积
        survival_area = {}
        growth_area = {}

        if not is_shortage:
            # 供水充足：全部存活且正常生长
            for crop in self.crops:
                survival_area[crop] = self.crops[crop]['area']
                growth_area[crop] = self.crops[crop]['area']
        else:
            # 供水不足：按比例缩减
            shortage_ratio = 1 - (supply_data['total_supply'] / self.daily_water_demand)
            for crop in self.crops:
                # 存活面积 = 种植面积 × (1 - 缺口比例×作物优先级系数)
                survival_area[crop] = max(0, self.crops[crop]['area'] * (1 - shortage_ratio * (1 - crop_ratios[crop])))
                # 正常生长面积为存活面积的60%（假设）
                growth_area[crop] = survival_area[crop] * 0.6

        # 生成表3结果
        table3 = pd.DataFrame({
            '作物': list(self.crops.keys()),
            '种植面积(公顷)': [self.crops[c]['area'] for c in self.crops],
            '存活面积(公顷)': [round(survival_area[c], 2) for c in self.crops],
            '正常生长面积(公顷)': [round(growth_area[c], 2) for c in self.crops]
        })

        print("\n表3 作物存活情况表：")
        print(table3)
        return table3

    def emergency_water_strategy(self):
        """分析应急储备比例与旱灾概率的关系"""
        # 1. 基于正态分布模型计算不同概率下的需储备量
        # 假设旱灾时额外需水量服从均值为储水罐容积50%的正态分布
        mu = self.tank_total_volume * 0.5
        sigma = mu * 0.3

        # 2. 计算各概率对应的储备比例
        table4_data = []
        for prob in self.drought_probabilities:
            # 旱灾概率p对应分位数z（右侧概率）
            z = norm.ppf(1 - prob / 100)
            required_volume = mu + z * sigma
            required_ratio = max(0.1,
                                 min(1.0, required_volume / self.tank_total_volume))  # 限制在10%-100%

            # 判断是否能保证全部存活
            is_fully_survive = required_ratio * self.tank_total_volume + \
                               self.pipe_daily_flow * 0.8 >= self.daily_water_demand

            table4_data.append({
                '旱灾概率(%)': prob,
                '建议应急储备水源比例(%)': round(required_ratio * 100, 1),
                '此比例能否保证作物全部存活': '是' if is_fully_survive else '否'
            })

        # 生成表4结果
        table4 = pd.DataFrame(table4_data)
        print("\n表4 应急水源比例与旱灾概率关系表：")
        print(table4)
        return table4

    def run_analysis(self):
        """执行问题3完整分析流程"""
        print("=== 问题3：旱灾应急分析 ===")
        # 步骤1：计算旱灾供水能力
        supply_data = self.calculate_drought_supply()

        # 步骤2：分析作物存活情况
        table3 = self.analyze_crop_survival(supply_data)

        # 步骤3：制定应急储备策略
        table4 = self.emergency_water_strategy()

        # 可视化应急储备比例与旱灾概率关系
        plt.figure(figsize=(10, 6))
        plt.plot(table4['旱灾概率(%)'], table4['建议应急储备水源比例(%)'], 'o-', color='red')
        plt.xlabel('旱灾概率(%)')
        plt.ylabel('建议应急储备水源比例(%)')
        plt.title('应急储备比例与旱灾概率关系')
        plt.grid(linestyle='--', alpha=0.7)
        plt.tight_layout()
        plt.show()

        return {
            '表3': table3,
            '表4': table4
        }


if __name__ == "__main__":
    # 这里假设已经获取到问题2的运行结果，以字典形式传入
    problem2_result = {
        '管道日流量': 3728.50,  # 这里填入问题2中实际计算出的管道日流量
        '储水罐容积': 18000.00,  # 这里填入问题2中实际计算出的储水罐容积
        '喷头位置': np.array([  # 这里填入问题2中实际计算出的喷头位置
            [25, 25], [25, 75], [75, 25], [75, 75]
        ]),
        '日均需水量': 33556.46 / 31  # 这里填入问题2中7月总需水量除以天数得到的日均需水量
    }
    analyzer = DroughtEmergencyAnalyzer(problem2_result)
    result = analyzer.run_analysis()